import cv2
import cupy as cp
import numpy as np
import os
from scipy.optimize import curve_fit
from scipy.spatial import cKDTree
from skimage.measure import label, regionprops
from skimage.morphology import skeletonize
import pandas as pd
import matplotlib.pyplot as plt
import warnings
from typing import List, Tuple, Dict, Any, Optional, Union
import time
import re
import multiprocessing
from concurrent.futures import ProcessPoolExecutor, as_completed
import gc
from scipy.interpolate import interp1d

warnings.filterwarnings('ignore')

# ==================== 全局配置 ====================
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

# 输出加速状态
print("=" * 60)
print("加速状态检查:")
print(f"CuPy 设备: {cp.cuda.runtime.getDeviceProperties(0)['name']}")
print("=" * 60)

# 钻孔参数（基于文档表3）
BOREHOLE_DIAMETER = 30.0  # mm
BOREHOLE_CIRCUMFERENCE = np.pi * BOREHOLE_DIAMETER  # ≈94.25 mm
BOREHOLE_DEPTH = 500.0  # mm (单孔深度)

# DPI参数（基于文档数据说明）
MULTI_BOREHOLE_DPI = 232.85  # 多孔图像DPI（附件4）
MM_PER_INCH = 25.4  # 1英寸=25.4毫米

# 输入输出路径
INPUT_DIR = "dataset/"  # 修改为数据集根目录
OUTPUT_DIR = "results/"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 采样方法配置
SAMPLING_METHODS = ['equidistant', 'arc_length', 'random', 'curvature_based']
DEFAULT_NUM_POINTS = 100

# 钻孔空间位置配置（基于文档表3）
BOREHOLE_CONFIG = {
    '1#': {'coords': (500, 2000, 0), 'depth': 7000},
    '2#': {'coords': (1500, 2000, 0), 'depth': 7000},
    '3#': {'coords': (2500, 2000, 0), 'depth': 7000},
    '4#': {'coords': (500, 1000, 0), 'depth': 5000},
    '5#': {'coords': (1500, 1000, 0), 'depth': 7000},
    '6#': {'coords': (2500, 1000, 0), 'depth': 7000}
}

# GPU配置
NUM_GPUS = 1
cp.cuda.Device(0).use()


# ==================== 核心函数定义 ====================

def sine_func(x: np.ndarray, R: float, P: float, beta: float, C: float) -> np.ndarray:
    """正弦函数模型：y = R * sin(2πx/P + beta) + C"""
    x = cp.asarray(x)
    result = R * cp.sin(2 * cp.pi * x / P + beta) + C
    return cp.asnumpy(result)  # 返回NumPy数组以兼容SciPy


def pixel_to_mm_coordinates(x_pixel: cp.ndarray, y_pixel: cp.ndarray,
                            image_width: int, image_height: int,
                            dpi: float = MULTI_BOREHOLE_DPI) -> Tuple[cp.ndarray, cp.ndarray]:
    """
    将像素坐标转换为毫米坐标（支持多孔DPI）- 使用CuPy加速，返回CuPy数组
    """
    mm_per_pixel = MM_PER_INCH / dpi
    x_mm = x_pixel * mm_per_pixel
    y_mm = y_pixel * mm_per_pixel
    return x_mm, y_mm


def enhance_image(image: np.ndarray) -> np.ndarray:
    """增强图像质量，减少噪声并连接断裂的裂隙 - 使用OpenCV CUDA加速"""
    # 检查CUDA是否可用
    if not cv2.cuda.getCudaEnabledDeviceCount():
        raise RuntimeError("CUDA不可用，请检查OpenCV CUDA模块")

    # 上传图像到GPU
    gpu_image = cv2.cuda_GpuMat()
    gpu_image.upload(image)

    # 使用CUDA版本的CLAHE
    clahe = cv2.cuda.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    gpu_enhanced = clahe.apply(gpu_image, cv2.cuda_Stream.Null())

    # 应用中值滤波去噪（CUDA版本）
    median_filter = cv2.cuda.createMedianFilter(cv2.CV_8UC1, 5)
    gpu_filtered = median_filter.apply(gpu_enhanced)

    # 使用形态学操作连接断裂区域（CUDA版本）
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    morph_filter = cv2.cuda.createMorphologyFilter(cv2.MORPH_CLOSE, cv2.CV_8UC1, kernel)
    gpu_closed = morph_filter.apply(gpu_filtered)

    # 下载回CPU
    result = gpu_closed.download()
    return result


def validate_fracture(coords: cp.ndarray, min_length_mm=30, min_x_span_mm=30.0,
                      min_y_std_mm=2.0) -> Tuple[bool, str]:
    """验证识别出的区域是否为有效的正弦状裂隙 - 使用CuPy加速计算，接受CuPy数组"""
    if len(coords) == 0:
        return False, "坐标数组为空"

    # 直接使用CuPy数组
    x_data, y_data = coords[:, 0], coords[:, 1]

    # 计算长度
    dx = cp.diff(x_data)
    dy = cp.diff(y_data)
    length = cp.sum(cp.sqrt(dx ** 2 + dy ** 2))

    # 计算x方向跨度
    x_span = cp.max(x_data) - cp.min(x_data)

    # 计算y方向标准差
    y_std = cp.std(y_data)

    # 将结果转换为标量
    length = float(length)
    x_span = float(x_span)
    y_std = float(y_std)

    if length < min_length_mm:
        return False, f"长度不足({length:.2f}mm < {min_length_mm}mm)"

    if x_span < min_x_span_mm:
        return False, f"x方向跨度不足({x_span:.2f}mm < {min_x_span_mm}mm)"

    if y_std < min_y_std_mm:
        return False, f"y方向变化太小({y_std:.2f}mm < {min_y_std_mm}mm)"

    return True, "验证通过"


def calculate_r_squared(y_true: cp.ndarray, y_pred: cp.ndarray) -> float:
    """计算决定系数R² - 使用CuPy加速计算，接受CuPy数组"""
    ss_res = cp.sum((y_true - y_pred) ** 2)
    ss_tot = cp.sum((y_true - cp.mean(y_true)) ** 2)
    r_squared = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0.0
    return float(r_squared)


def robust_curve_fit(x_data: np.ndarray, y_data: np.ndarray,
                     R_init: float, P_init: float, beta_init: float, C_init: float) -> tuple:
    """鲁棒的曲线拟合 - 多策略优化以避免局部最优解"""
    # 如果数据点过多，进行下采样以避免内存问题
    if len(x_data) > 10000:
        indices = np.random.choice(len(x_data), 10000, replace=False)
        x_data = x_data[indices]
        y_data = y_data[indices]

    # 将数据传回CPU，因为curve_fit需要NumPy数组
    x_data = cp.asnumpy(x_data)
    y_data = cp.asnumpy(y_data)

    # 设置合理的参数边界
    R_min, R_max = max(0.1, R_init * 0.1), min(100.0, R_init * 10.0)
    P_min, P_max = max(10.0, P_init * 0.1), min(1000.0, P_init * 10.0)
    C_min, C_max = np.min(y_data) - 50, np.max(y_data) + 50

    # 确保边界有效
    if R_min >= R_max:
        R_min, R_max = 0.1, 100.0
    if P_min >= P_max:
        P_min, P_max = 10.0, 1000.0
    if C_min >= C_max:
        C_min, C_max = np.min(y_data) - 50, np.max(y_data) + 50

    bounds = ([R_min, P_min, -2 * np.pi, C_min], [R_max, P_max, 2 * np.pi, C_max])

    # 裁剪初始值到边界内
    R_init_clipped = np.clip(R_init, R_min, R_max)
    P_init_clipped = np.clip(P_init, P_min, P_max)
    C_init_clipped = np.clip(C_init, C_min, C_max)
    beta_init_clipped = np.clip(beta_init, -2 * np.pi, 2 * np.pi)
    p0 = [R_init_clipped, P_init_clipped, beta_init_clipped, C_init_clipped]

    # 使用Levenberg-Marquardt算法进行拟合
    popt, _ = curve_fit(sine_func, x_data, y_data, p0=p0, bounds=bounds,
                        method='trf', max_nfev=10000, ftol=1e-6, xtol=1e-6)

    R, P, beta, C = popt
    beta = (beta + np.pi) % (2 * np.pi) - np.pi  # 规范化beta到[-π, π]

    # 计算R²
    y_pred = sine_func(x_data, R, P, beta, C)
    r_squared = calculate_r_squared(y_data, y_pred)

    return (R, P, beta, C, r_squared)


def extract_contour_skeleton(region_mask: np.ndarray) -> np.ndarray:
    """使用骨架化算法提取裂隙区域的轮廓线"""
    if region_mask.dtype != np.uint8:
        region_mask = region_mask.astype(np.uint8)

    # 使用更高效的骨架化算法
    skeleton = skeletonize(region_mask > 0, method='lee')
    y_coords, x_coords = np.where(skeleton)

    # 如果点数过多，进行下采样
    if len(x_coords) > 1000:
        indices = np.random.choice(len(x_coords), 1000, replace=False)
        x_coords, y_coords = x_coords[indices], y_coords[indices]

    contour_points = np.column_stack((x_coords, y_coords))
    return contour_points if len(contour_points) >= 2 else np.array([])


def calculate_jrc(x_points: cp.ndarray, y_points: cp.ndarray,
                  sampling_method: str = 'equidistant', num_points: int = DEFAULT_NUM_POINTS) -> Tuple[float, float]:
    """计算JRC值基于离散坐标点，支持四种采样方法，接受CuPy数组"""
    # 转换为NumPy数组用于计算
    x_points_np = cp.asnumpy(x_points)
    y_points_np = cp.asnumpy(y_points)

    # 去除重复点并按x排序
    points = np.column_stack((x_points_np, y_points_np))
    unique_points = np.unique(points, axis=0)
    if len(unique_points) < 2:
        return 0.0, 0.0

    x_sorted = unique_points[:, 0]
    y_sorted = unique_points[:, 1]
    sorted_indices = np.argsort(x_sorted)
    x_sorted, y_sorted = x_sorted[sorted_indices], y_sorted[sorted_indices]

    if sampling_method == 'equidistant':
        # 等距离采样
        x_interp = np.linspace(np.min(x_sorted), np.max(x_sorted), num_points)
        f = interp1d(x_sorted, y_sorted, kind='linear', fill_value='extrapolate')
        y_interp = f(x_interp)

    elif sampling_method == 'arc_length':
        # 弧长采样
        dx = np.diff(x_sorted)
        dy = np.diff(y_sorted)
        ds = np.sqrt(dx ** 2 + dy ** 2)
        s = np.concatenate(([0], np.cumsum(ds)))
        total_arc_length = s[-1]

        if total_arc_length == 0:
            x_interp = np.linspace(np.min(x_sorted), np.max(x_sorted), num_points)
            f = interp1d(x_sorted, y_sorted, kind='linear', fill_value='extrapolate')
            y_interp = f(x_interp)
        else:
            s_interp = np.linspace(0, total_arc_length, num_points)
            f_x = interp1d(s, x_sorted, kind='linear', fill_value='extrapolate')
            f_y = interp1d(s, y_sorted, kind='linear', fill_value='extrapolate')
            x_interp = f_x(s_interp)
            y_interp = f_y(s_interp)

    elif sampling_method == 'random':
        # 随机采样
        if len(x_sorted) < num_points:
            indices = np.random.choice(len(x_sorted), size=num_points, replace=True)
        else:
            indices = np.random.choice(len(x_sorted), size=num_points, replace=False)

        x_interp = x_sorted[indices]
        y_interp = y_sorted[indices]
        sort_idx = np.argsort(x_interp)
        x_interp, y_interp = x_interp[sort_idx], y_interp[sort_idx]

    elif sampling_method == 'curvature_based':
        # 曲率采样
        dx = np.diff(x_sorted)
        dy = np.diff(y_sorted)
        dx = np.where(dx == 0, 1e-5, dx)  # 避免除零

        # 计算一阶导数
        derivatives = dy / dx

        # 计算二阶导数（曲率）
        d2x = np.diff(x_sorted[:-1])
        d2y = np.diff(derivatives)
        d2x = np.where(d2x == 0, 1e-5, d2x)  # 避免除零

        curvature = np.abs(d2y / d2x)
        curvature = np.concatenate(([0], curvature, [0]))  # 补齐长度

        # 归一化权重
        weights = curvature / np.sum(curvature) if np.sum(curvature) > 0 else np.ones_like(x_sorted) / len(x_sorted)

        # 根据曲率权重采样
        indices = np.random.choice(len(x_sorted), size=num_points, p=weights, replace=True)
        x_interp = x_sorted[indices]
        y_interp = y_sorted[indices]
        sort_idx = np.argsort(x_interp)
        x_interp, y_interp = x_interp[sort_idx], y_interp[sort_idx]

    else:
        raise ValueError(f"不支持的采样方法: {sampling_method}")

    # 计算JRC值
    dx = np.diff(x_interp)
    dy = np.diff(y_interp)
    dx = np.where(dx == 0, 1e-5, dx)  # 避免除零

    derivatives_squared = (dy / dx) ** 2
    z2 = np.sqrt(np.mean(derivatives_squared))
    jrc = 51.85 * (z2 ** 0.6) - 10.37

    return max(0.0, jrc), z2


def create_visualization(enhanced_image: np.ndarray, fitted_curves: list,
                         plot_output: str, image_width: int, image_height: int) -> None:
    """创建对比可视化图像"""
    image_width_mm = (image_width * MM_PER_INCH) / MULTI_BOREHOLE_DPI
    image_height_mm = (image_height * MM_PER_INCH) / MULTI_BOREHOLE_DPI

    fig, ax = plt.subplots(figsize=(16, 10))
    ax.imshow(enhanced_image, cmap='gray', extent=[0, image_width_mm, image_height_mm, 0], alpha=0.2)

    # 使用更丰富的极颜色方案
    fracture_colors = plt.cm.Set3(np.linspace(0, 1, len(fitted_curves)))

    for i, (x_data, y_data, _) in enumerate(fitted_curves):
        ax.scatter(x_data, y_data, color=fracture_colors[i], s=5, alpha=0.8, edgecolors='w', linewidth=0.5)

    for i, (x_data, _, params) in enumerate(fitted_curves):
        x_plot = np.linspace(min(x_data), max(x_data), 200)
        y_plot = sine_func(x_plot, *params)
        ax.plot(x_plot, y_plot, color=fracture_colors[i], linewidth=2.0, alpha=0.9)


    ax.set_xlim(0, image_width_mm)
    ax.set_ylim(0, image_height_mm)
    ax.set_xlabel('周向距离 (mm)', fontsize=12)
    ax.set_ylabel('轴向深度 (mm)', fontsize=12)
    ax.set_title('裂隙识别与拟合结果对比', fontsize=14, fontweight='bold')
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(plot_output, dpi=300, bbox_inches='tight')
    plt.close()


def convert_to_global_coordinates(contour_points: np.ndarray, borehole_id: str) -> np.ndarray:
    """将轮廓点转换到全局三维坐标系"""
    borehole_info = BOREHOLE_CONFIG[borehole_id]
    x_base, y_base, z_base = borehole_info['coords']

    global_points = []
    for point in contour_points:
        x_mm, y_mm = point[0], point[1]

        # 将周向距离转换为角度
        theta = 2 * np.pi * x_mm / BOREHOLE_CIRCUMFERENCE

        # 计算全局坐标
        x_global = x_base + np.cos(theta) * (BOREHOLE_DIAMETER / 2)
        y_global = y_base + np.sin(theta) * (BOREHOLE_DIAMETER / 2)
        z_global = z_base + y_mm  # y_mm是深度方向

        global_points.append([x_global, y_global, z_global])

    return np.array(global_points)


def precompute_kdtrees(fractures):
    """预先计算所有裂隙的KDTree"""
    kdtrees = []
    for fracture in fractures:
        points = fracture['global_points']
        kdtree = cKDTree(points)
        kdtrees.append(kdtree)
    return kdtrees


def calculate_connectivity_probability(fracture1: Dict[str, Any], fracture2: Dict[str, Any],
                                       kdtree1: cKDTree = None, kdtree2: cKDTree = None) -> float:
    """计算两条裂隙之间的连通概率 - 使用预先计算的KDTree"""
    points1 = fracture1['global_points']
    points2 = fracture2['global_points']

    # 使用预先计算的KDTree或创建新的
    if kdtree1 is None:
        kdtree1 = cKDTree(points1)
    if kdtree2 is None:
        kdtree2 = cKDTree(points2)

    # 计算最近距离
    dist1, _ = kdtree1.query(points2)
    dist2, _ = kdtree2.query(points1)
    min_distance = min(np.min(dist1), np.min(dist2))

    # 距离衰减函数
    distance_score = np.exp(-min_distance / 500.0)

    # 方向一致性
    params1 = fracture1['params']
    params2 = fracture2['params']

    direction1 = np.array([params1['R'], params1['P'], params1['beta']])
    direction2 = np.array([params2['P'], params2['P'], params2['beta']])

    norm1 = np.linalg.norm(direction1)
    norm2 = np.linalg.norm(direction2)

    if norm1 > 0 and norm2 > 0:
        direction_similarity = np.dot(direction1, direction2) / (norm1 * norm2)
        direction_score = max(0, direction_similarity)
    else:
        direction_score = 0.0

    # 粗糙度因素
    jrc1 = fracture1['jrc_values']['equidistant']
    jrc2 = fracture2['jrc_values']['equidistant']
    roughness_factor = 0.5 + 0.5 * np.tanh((jrc1 + jrc2 - 15) / 10)

    # 综合概率
    connectivity_prob = (0.4 * distance_score +
                         0.3 * direction_score +
                         0.3 * roughness_factor)

    # 确保概率在[0,1]范围内
    connectivity_prob = np.clip(connectivity_prob, 0.0, 1.0)

    return float(connectivity_prob)


def compute_connectivity_probability_with_kdtree(fractures, kdtrees, i, j):
    """使用预先计算的KDTree计算两条裂隙之间的连通概率（用于多进程）"""
    fracture1 = fractures[i]
    fracture2 = fractures[j]
    kdtree1 = kdtrees[i]
    kdtree2 = kdtrees[j]

    return calculate_connectivity_probability(fracture1, fracture2, kdtree1, kdtree2)


def calculate_connectivity_matrix(fractures: List[Dict[str, Any]]) -> np.ndarray:
    """优化连通性概率矩阵计算，使用预先计算的KDTree"""
    n = len(fractures)
    connectivity_matrix = np.zeros((n, n))
    print("计算连通性概率矩阵...")
    start_time = time.time()

    # 预先计算所有裂隙的KDTree
    kdtrees = precompute_kdtrees(fractures)

    # 仅处理不同钻孔的裂隙对
    borehole_groups = {}
    for i, fracture in enumerate(fractures):
        borehole_id = fracture['borehole_id']
        if borehole_id not in borehole_groups:
            borehole_groups[borehole_id] = []
        borehole_groups[borehole_id].append(i)

    # 使用多进程并行计算连通性
    with ProcessPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
        futures = {}
        for src_borehole, src_indices in borehole_groups.items():
            for tgt_borehole, tgt_indices in borehole_groups.items():
                if src_borehole == tgt_borehole:
                    continue

                for i in src_indices:
                    for j in tgt_indices:
                        if j > i:  # 避免重复计算
                            future = executor.submit(
                                compute_connectivity_probability_with_kdtree,
                                fractures, kdtrees, i, j
                            )
                            futures[future] = (i, j)

        for future in as_completed(futures):
            i, j = futures[future]
            prob = future.result()
            connectivity_matrix[i, j] = prob
            connectivity_matrix[j, i] = prob

    elapsed_time = time.time() - start_time
    print(f"连通性矩阵计算完成，耗时 {elapsed_time:.2f}秒")
    return connectivity_matrix


def generate_3d_visualization(fractures: List[Dict[str, Any]],
                              connectivity_matrix: np.ndarray,
                              output_path: str) -> None:
    """优化三维裂隙网络可视化，使用KDTree加速最近邻搜索"""
    print("生成三维可视化...")
    fig = plt.figure(figsize=(16, 12))
    ax = fig.add_subplot(111, projection='3d')

    # 绘制钻孔轨迹
    for borehole_id, config in BOREHOLE_CONFIG.items():
        x, y, z_start = config['coords']
        z_end = z_start + config['depth']
        ax.plot([x, x], [y, y], [z_start, z_end],
                'k-', linewidth=3, label=f'钻孔 {borehole_id}')

    # 优化点云绘制 - 使用更高效的标记方法
    borehole_colors = {
        '1#': '#FF6B6B', '2#': '#4ECDC4', '3#': '#45B7D1',
        '4#': '#FFD166', '5#': '#118AB2', '6#': '#9B5DE5'
    }

    # 按钻孔分组点云
    grouped_points = {bid: [] for bid in borehole_colors.keys()}
    for fracture in fractures:
        bid = fracture['borehole_id']
        if bid in grouped_points:
            grouped_points[bid].append(fracture['global_points'])

    # 批量绘制点
    for bid, points_list in grouped_points.items():
        if not points_list:
            continue
        all_points = np.vstack(points_list)
        # 下采样以减少点数
        if len(all_points) > 10000:
            indices = np.random.choice(len(all_points), 10000, replace=False)
            all_points = all_points[indices]
        ax.scatter(all_points[:, 0], all_points[:, 1], all_points[:, 2],
                   c=borehole_colors[bid], s=10, alpha=0.7, label=bid)

    # 优化连通路径绘制
    connection_count = 0
    max_connections = 500  # 限制最大连接数

    # 预先为所有裂隙构建KDTree
    kdtrees = precompute_kdtrees(fractures)

    # 收集所有需要绘制的连接
    connections = []
    for i in range(len(fractures)):
        for j in range(i + 1, len(fractures)):
            if connectivity_matrix[i, j] > 0.7 and connection_count < max_connections:
                connections.append((i, j))
                connection_count += 1

    # 批量处理连接
    for i, j in connections:
        points_i = fractures[i]['global_points']
        points_j = fractures[j]['global_points']

        # 使用KDTree查找最近点对
        tree_i = kdtrees[i]
        tree_j = kdtrees[j]

        # 查询points_j在points_i中的最近邻
        dist_i, idx_i = tree_i.query(points_j)
        min_idx_j = np.argmin(dist_i)
        min_idx_i = idx_i[min_idx_j]

        # 绘制连线
        prob = connectivity_matrix[i, j]
        ax.plot([points_i[min_idx_i, 0], points_j[min_idx_j, 0]],
                [points_i[min_idx_i, 1], points_j[min_idx_j, 1]],
                [points_i[min_idx_i, 2], points_j[min_idx_j, 2]],
                'r-', alpha=prob, linewidth=2 * prob)

    print(f"绘制了 {len(connections)} 条高概率连通路径")

    ax.set_xlabel('X (mm)', fontsize=12, fontweight='bold')
    ax.set_ylabel('Y (mm)', fontsize=12, fontweight='bold')
    ax.set_zlabel('Z (mm)', fontsize=12, fontweight='bold')
    ax.set_title('多钻孔裂隙网络三维可视化', fontsize=16, fontweight='bold')
    ax.view_init(elev=30, azim=45)

    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    ax.legend(by_label.values(), by_label.keys(), loc='upper left', bbox_to_anchor=(0, 1))

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"三维可视化已保存至: {output_path}")


def save_connectivity_results(connectivity_matrix: np.ndarray,
                              fractures,
                              output_path: str) -> None:
    """
    保存连通性分析结果到CSV文件
    格式符合文档中表1的要求，使用KDTree加速最近邻搜索
    """
    results = []

    # 预先为所有裂隙构建KDTree
    kdtrees = precompute_kdtrees(fractures)

    for i in range(len(fractures)):
        for j in range(i + 1, len(fractures)):
            if fractures[i]['borehole_id'] != fractures[j]['borehole_id'] and connectivity_matrix[i, j] > 0:
                points_i = fractures[i]['global_points']
                points_j = fractures[j]['global_points']

                # 使用KDTree计算最近距离
                tree_i = kdtrees[i]
                tree_j = kdtrees[j]

                dist_i, _ = tree_i.query(points_j)
                dist_j, _ = tree_j.query(points_i)
                min_distance = min(np.min(dist_i), np.min(dist_j))

                results.append({
                    'Source_Borehole': fractures[i]['borehole_id'],
                    'Target_Borehole': fractures[j]['borehole_id'],
                    'Connectivity_Probability': round(connectivity_matrix[i, j], 4),
                    'Source_Amplitude_R': round(fractures[i]['params']['R'], 2),
                    'Target_Amplitude_R': round(fractures[j]['params']['R'], 2),
                    'Source_Period_P': round(fractures[i]['params']['P'], 2),
                    'Target_Period_P': round(fractures[j]['params']['P'], 2),
                    'Source_JRC': round(fractures[i]['jrc_values']['equidistant'], 2),
                    'Target_JRC': round(fractures[j]['jrc_values']['equidistant'], 2),
                    'Distance_mm': round(min_distance, 2)
                })

    df = pd.DataFrame(results)
    df.to_csv(output_path, index=False, encoding='utf-8-sig')
    print(f"连通性结果已保存至: {output_path}")


def analyze_connectivity_uncertainty(connectivity_matrix: np.ndarray,
                                     fractures: List[Dict[str, Any]]) -> List[Tuple[float, float, float]]:
    """
    分析连通性不确定性并推荐补充钻孔位置
    返回按优先级排序的钻孔位置列表
    """
    print("分析连通性不确定性...")

    # 计算每个裂隙的不确定性（连通概率方差）
    uncertainty_scores = {}
    for i in range(len(fractures)):
        probs = connectivity_matrix[i]
        # 只考虑非零概率
        valid_probs = probs[probs > 0]
        if len(valid_probs) > 0:
            uncertainty = np.var(valid_probs)
        else:
            uncertainty = 0
        uncertainty_scores[i] = uncertainty

    # 识别高不确定性区域（前3个）
    high_uncertainty_regions = sorted(
        uncertainty_scores.items(),
        key=lambda x: x[1],
        reverse=True
    )[:3]

    # 基于高不确定性区域推荐钻孔位置
    recommended_locations = []
    for region_idx, uncertainty in high_uncertainty_regions:
        fracture = fractures[region_idx]
        points = fracture['global_points']

        # 计算区域中心点
        center_x = np.mean(points[:, 0])
        center_y = np.mean(points[:, 1])
        center_z = np.mean(points[:, 2])

        # 推荐钻孔位置（基于文档中的坐标系）
        recommended_locations.append((center_x, center_y, center_z))

    # 如果推荐位置不足3个，添加默认位置
    while len(recommended_locations) < 3:
        # 基于文档中的钻孔阵列布局推荐
        default_locations = [
            (1000, 1500, 0),  # 阵列中心
            (2000, 1500, 0),  # 东部区域
            (1500, 1000, 0)  # 南部区域
        ]
        recommended_locations.append(default_locations[len(recommended_locations)])

    print("不确定性分析完成")
    return recommended_locations


def process_single_borehole_image(image_path: str, borehole_id: str, depth_offset: float = 0.0, gpu_id: int = 0) -> \
        List[Dict[str, Any]]:
    """处理单个钻孔图像，返回裂隙特征列表 - 支持GPU设备分配"""
    # 设置CuPy设备
    cp.cuda.Device(gpu_id).use()

    # 读取图像 - 添加检查确保图像存在
    if not os.path.exists(image_path):
        print(f"图像文件不存在: {image_path}")
        return []

    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if image is None:
        print(f"无法读取图像: {image_path}")
        return []

    height, width = image.shape

    # 增强图像（使用OpenCV CUDA）
    enhanced_image = enhance_image(image)

    # 标记连通区域
    labeled_image = label(enhanced_image, connectivity=2)
    regions = regionprops(labeled_image)

    fractures = []
    for region in regions:
        # 提取区域坐标
        y_pix, x_pix = region.coords[:, 0], region.coords[:, 1]

        # 转换为CuPy数组
        x_pix_gpu = cp.asarray(x_pix)
        y_pix_gpu = cp.asarray(y_pix)

        # 转换为物理坐标（使用CuPy加速，返回CuPy数组）
        x_phys, y_phys = pixel_to_mm_coordinates(x_pix_gpu, y_pix_gpu, width, height, MULTI_BOREHOLE_DPI)

        # 应用深度偏移
        y_phys = y_phys + depth_offset

        # 构建坐标数组（CuPy）
        region_coords = cp.column_stack((x_phys, y_phys))

        # 验证是否为有效裂隙（直接传递CuPy数组）
        is_valid, reason = validate_fracture(region_coords)
        if is_valid:
            fractures.append(region_coords)

    results = []
    for i, coords in enumerate(fractures):
        x_data, y_data = coords[:, 0], coords[:, 1]  # CuPy数组

        # 转换为NumPy数组用于曲线拟合（因为SciPy需要NumPy数组）
        x_data_np = cp.asnumpy(x_data)
        y_data_np = cp.asnumpy(y_data)

        # 设置初始极参数
        P_init = BOREHOLE_CIRCUMFERENCE
        C_init = np.median(y_data_np)
        R_init = (np.max(y_data_np) - np.min(y_data_np)) / 2
        beta_init = 0.0

        # 曲线拟合
        R, P, beta, C, r_squared = robust_curve_fit(x_data_np, y_data_np, R_init, P_init, beta_init, C_init)

        if r_squared >= 0.3:  # 拟合优度阈值
            # 提取轮廓骨架（使用CPU算法）
            region_mask = (labeled_image == region.label).astype(np.uint8)
            contour_points_pixel = extract_contour_skeleton(region_mask)

            if len(contour_points_pixel) > 0:
                # 转换为CuPy数组
                x_contour_pix = cp.asarray(contour_points_pixel[:, 0])
                y_contour_pix = cp.asarray(contour_points_pixel[:, 1])

                # 转换为物理坐标（返回CuPy数组）
                x_contour, y_contour = pixel_to_mm_coordinates(
                    x_contour_pix, y_contour_pix, width, height, MULTI_BOREHOLE_DPI
                )

                # 应用深度偏移
                y_contour = y_contour + depth_offset

                # 计算不同采样方法的JRC值（传递CuPy数组）
                jrc_results = {}
                for method in SAMPLING_METHODS:
                    jrc_value, _ = calculate_jrc(x_contour, y_contour, sampling_method=method)
                    jrc_results[method] = jrc_value
            else:
                jrc_results = {method: 0.0 for method in SAMPLING_METHODS}

            # 转换为全局坐标（需要NumPy数组，所以转换）
            coords_np = cp.asnumpy(coords)
            global_points = convert_to_global_coordinates(coords_np, borehole_id)

            results.append({
                'borehole_id': borehole_id,
                'global_points': global_points,
                'params': {'R': R, 'P': P, 'beta': beta, 'C': C},
                'jrc_values': jrc_results,
                'r_squared': r_squared
            })

    # 释放内存
    del labeled_image
    del regions
    gc.collect()

    return results


def process_multiborehole_data(input_dir: str, output_dir: str) -> List[Dict[str, Any]]:
    """处理多钻孔数据，使用多进程并行处理 - 分配GPU设备"""
    print("开始处理多钻孔数据...")
    start_time = time.time()
    all_fractures = []
    image_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.tiff']

    # 收集所有需要处理的任务
    tasks = []
    for borehole_id in BOREHOLE_CONFIG.keys():
        borehole_folder = os.path.join(input_dir, f"hole_{borehole_id.replace('#', '')}")
        if not os.path.exists(borehole_folder):
            print(f"警告：钻孔 {borehole_id} 的文件夹不存在: {borehole_folder}")
            continue

        label_result_dir = os.path.join(borehole_folder, "label_result")
        if not os.path.exists(label_result_dir):
            print(f"警告：钻孔 {borehole_id} 极标签结果文件夹不存在: {label_result_dir}")
            continue

        mask_files = [f for f in os.listdir(label_result_dir)
                      if any(f.lower().endswith(ext) for ext in image_extensions)]
        mask_files.sort(key=lambda x: [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', x)])

        if not mask_files:
            print(f"警告：钻孔 {borehole_id} 的标签结果文件夹中没有找到图像文件")
            continue

        depth_interval = BOREHOLE_CONFIG[borehole_id]['depth'] / len(mask_files)

        for i, mask_file in enumerate(mask_files):
            image_path = os.path.join(label_result_dir, mask_file)
            depth_offset = i * depth_interval
            # 分配GPU ID循环使用
            gpu_id = i % NUM_GPUS
            tasks.append((image_path, borehole_id, depth_offset, gpu_id))

    # 使用多进程并行处理
    with ProcessPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
        futures = [executor.submit(process_single_borehole_image, *task) for task in tasks]
        for future in as_completed(futures):
            fractures = future.result()
            if fractures:
                all_fractures.extend(fractures)

    elapsed_time = time.time() - start_time
    print(f"多钻孔处理完成，总共发现 {len(all_fractures)} 条裂隙，耗时 {elapsed_time:.2f} 秒")
    return all_fractures


def problem4_main():
    """
    问题4主函数：多钻孔裂隙网络的连通性分析与三维重构
    整合所有功能，完成从数据加载到结果输出的完整流程
    """
    print("=" * 80)
    print("开始问题4：多钻孔裂隙网络的连通性分析与三维重构")
    print("=" * 80)

    # 1. 处理多钻孔数据
    all_fractures = process_multiborehole_data(INPUT_DIR, OUTPUT_DIR)

    if not all_fractures:
        print("未发现有效裂隙，终止处理")
        return

    # 2. 计算连通性矩阵
    connectivity_matrix = calculate_connectivity_matrix(all_fractures)

    # 强制垃圾回收释放内存
    gc.collect()

    # 3. 生成三维可视化
    visualization_path = os.path.join(OUTPUT_DIR, "problem4_3d_network.png")
    generate_3d_visualization(all_fractures, connectivity_matrix, visualization_path)

    # 4. 保存连通性结果
    results_path = os.path.join(OUTPUT_DIR, "problem4_connectivity_results.csv")
    save_connectivity_results(connectivity_matrix, all_fractures, results_path)

    # 5. 不确定性分析和钻孔推荐
    recommended_drills = analyze_connectivity_uncertainty(connectivity_matrix, all_fractures)

    # 6. 输出推荐钻孔位置
    print("\n推荐补充钻孔位置（按优先级排序）:")
    for i, (x, y, z) in enumerate(recommended_drills, 1):
        print(f"优先级{i}: 坐标({x:.0f}, {y:.0f}, {z:.0f})")

    # 7. 保存推荐结果
    drill_df = pd.DataFrame(recommended_drills, columns=['X', 'Y', 'Z'])
    drill_df.to_csv(os.path.join(OUTPUT_DIR, "problem4_recommended_drills.csv"), index=False)

    print("=" * 80)
    print("问题4处理完成！")
    print("=" * 80)


# ==================== 主程序入口 ====================

if __name__ == "__main__":
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 运行问题4主函数
    problem4_main()